import re
import os
import time
from nltk.tokenize import word_tokenize
from nltk.corpus import stopwords
from nltk.stem import WordNetLemmatizer
from rich.progress import Progress, TextColumn, BarColumn, TimeElapsedColumn
from data_processing_common import sanitize_filename  # 导入 sanitize_filename

def get_text_from_generator(generator):
    """从生成器响应中提取文本"""
    response_text = ""
    try:
        while True:
            response = next(generator)
            choices = response.get('choices', [])
            for choice in choices:
                delta = choice.get('delta', {})
                if 'content' in delta:
                    response_text += delta['content']
    except StopIteration:
        pass
    return response_text

def process_single_image(image_path, image_inference, text_inference, silent=False, log_file=None):
    """处理单个图像文件生成元数据"""
    start_time = time.time()

    # 为此文件创建进度条实例
    with Progress(
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TimeElapsedColumn()
    ) as progress:
        task_id = progress.add_task(f"处理 {os.path.basename(image_path)}", total=1.0)
        foldername, filename, description = generate_image_metadata(image_path, progress, task_id, image_inference, text_inference)
    
    end_time = time.time()
    time_taken = end_time - start_time

    message = f"文件: {image_path}\n耗时: {time_taken:.2f} 秒\n描述: {description}\n文件夹名: {foldername}\n生成文件名: {filename}\n"
    if silent:
        if log_file:
            with open(log_file, 'a', encoding='utf-8') as f:
                f.write(message + '\n')
    else:
        print(message)
    return {
        'file_path': image_path,
        'foldername': foldername,
        'filename': filename,
        'description': description
    }

def process_image_files(image_paths, image_inference, text_inference, silent=False, log_file=None):
    """顺序处理多个图像文件"""
    data_list = []
    for image_path in image_paths:
        data = process_single_image(image_path, image_inference, text_inference, silent=silent, log_file=log_file)
        data_list.append(data)
    return data_list

def generate_image_metadata(image_path, progress, task_id, image_inference, text_inference):
    """为图像文件生成描述、文件夹名和文件名"""

    # 处理图像的总步骤数
    total_steps = 3

    # 步骤1: 使用image_inference生成描述
    description_prompt = "请详细描述这张图片，重点关注主体和重要细节。"
    description_generator = image_inference._chat(description_prompt, image_path)
    description = get_text_from_generator(description_generator).strip()
    progress.update(task_id, advance=1 / total_steps)

    # 步骤2: 使用text_inference生成文件名
    filename_prompt = f"""根据以下描述，为图像生成一个具体且描述性的文件名。
文件名最多使用3个词。使用名词，避免使用"描绘"、"显示"、"呈现"等动词开头。
不要包含"图片"、"jpg"、"png"等数据类型词。只使用字母并用下划线连接单词。

描述: {description}

示例:
描述: 一张山脉日落的照片。
文件名: 日落_山脉

现在请生成文件名。

只输出文件名，不要包含任何其他文本。

文件名:"""
    filename_response = text_inference.create_completion(filename_prompt)
    filename = filename_response['choices'][0]['text'].strip()
    # 如果存在"文件名:"前缀则移除
    filename = re.sub(r'^文件名:\s*', '', filename, flags=re.IGNORECASE).strip()
    progress.update(task_id, advance=1 / total_steps)

    # 步骤3: 使用text_inference从描述生成文件夹名
    foldername_prompt = f"""根据以下描述，生成最能代表此图像主题的通用类别。
这将用作文件夹名。类别最多使用2个词。使用名词，避免使用动词。
不要包含具体细节、文件名中的词或"未命名"、"未知"等通用术语。

描述: {description}

示例:
1. 描述: 一张山脉日落的照片。
   类别: 风景

2. 描述: 显示存储应用程序的智能手机图像，带有各种图标和信息。
   类别: 科技

3. 描述: 带有露珠的盛开红玫瑰特写。
   类别: 自然

现在请生成类别。

只输出类别，不要包含任何其他文本。

类别:"""
    foldername_response = text_inference.create_completion(foldername_prompt)
    foldername = foldername_response['choices'][0]['text'].strip()
    # 如果存在"类别:"前缀则移除
    foldername = re.sub(r'^类别:\s*', '', foldername, flags=re.IGNORECASE).strip()
    progress.update(task_id, advance=1 / total_steps)

    # 移除不需要的单词和停用词
    unwanted_words = set([
        '的', '和', '基于', '生成', '这个', '是', '文件名', '文件', '图像', '图片', '照片',
        '文件夹', '类别', '输出', '仅', '下面', '文本', 'jpg', 'png', 'jpeg', 'gif', 'bmp', 'svg',
        '标志', '在', '上', '的', '与', '由', '为', '到', '从', '一个', '作为', '在', '红色', '蓝色',
        '绿色', '颜色', '彩色', '文本', '图形', '主要', '主题', '重要',
        '细节', '描述', '描绘', '显示', '展示', '呈现', '说明', '特征',
        '提供', '覆盖', '包括', '演示', '描述'
    ])
    stop_words = set(stopwords.words('english'))
    all_unwanted_words = unwanted_words.union(stop_words)
    lemmatizer = WordNetLemmatizer()

    # 清理和处理AI输出的函数
    def clean_ai_output(text, max_words):
        # 移除文件扩展名和特殊字符
        text = re.sub(r'\.\w{1,4}$', '', text)  # 移除.jpg, .png等扩展名
        text = re.sub(r'[^\w\s]', ' ', text)  # 移除特殊字符
        text = re.sub(r'\d+', '', text)  # 移除数字
        text = text.strip()
        # 拆分连接词(如'GoogleChrome' -> 'Google Chrome')
        text = re.sub(r'([a-z])([A-Z])', r'\1 \2', text)
        # 分词和词形还原
        words = word_tokenize(text)
        words = [word.lower() for word in words if word.isalpha()]
        words = [lemmatizer.lemmatize(word) for word in words]
        # 移除不需要的词和重复词
        filtered_words = []
        seen = set()
        for word in words:
            if word not in all_unwanted_words and word not in seen:
                filtered_words.append(word)
                seen.add(word)
        # 限制最大词数
        filtered_words = filtered_words[:max_words]
        return '_'.join(filtered_words)

    # 处理文件名
    filename = clean_ai_output(filename, max_words=3)
    if not filename or filename.lower() in ('未命名', ''):
        # 使用描述中的关键词
        filename = clean_ai_output(description, max_words=3)
    if not filename:
        filename = '图片_' + os.path.splitext(os.path.basename(image_path))[0]

    sanitized_filename = sanitize_filename(filename, max_words=3)

    # 处理文件夹名
    foldername = clean_ai_output(foldername, max_words=2)
    if not foldername or foldername.lower() in ('未命名', ''):
        # 尝试从描述中提取关键词
        foldername = clean_ai_output(description, max_words=2)
        if not foldername:
            foldername = '图片'

    sanitized_foldername = sanitize_filename(foldername, max_words=2)

    return sanitized_foldername, sanitized_filename, description